import logging
import os

from llama_index.core import SummaryIndex
from llama_index.core.node_parser import SimpleNodeParser
from llama_index.llms.dashscope import DashScope
# 导入依赖 pip install llama-index-readers-wikipedia
from llama_index.readers.wikipedia import WikipediaReader

logging.basicConfig(level=logging.DEBUG)

print('初始化模型对象')
llm = DashScope(
    model="qwen-turbo",  # 使用通义千问Max模型 qwen-max
    api_key=os.getenv("DASHSCOPE_API_KEY"),  # 从阿里云控制台获取
    temperature=0.3
)

# 初始化数据加载器
loader = WikipediaReader()
# 创建文档
documents = loader.load_data(pages=["Shanghai"])
# 初始化节点解析器
parser = SimpleNodeParser.from_defaults()
# 获取节点
nodes = parser.get_nodes_from_documents(documents)
# 创建索引
index = SummaryIndex(nodes)
# 创建查询引擎
query_engine = index.as_query_engine(llm=llm)
print("对上海进行提问")
question = "Where is Shanghai?"
response = query_engine.query(question)
print('提问结果:', response)
